from kafka import KafkaProducer, KafkaConsumer, TopicPartition
from kafka.errors import KafkaError
import json
from typing import List, Dict, Any, Generator
from retrying import retry
from starmerx_warehouse.utils import log_utils

# 设置日志记录
logger = log_utils.log('kafka')


class KafkaTool:
    """
    Kafka工具类：
        1. 创建生产者对象
            1.1 生产消息
        2. 创建消费者对象
            2.1 消费消息
        3. 获取某主题下的消息总数
    """

    def __init__(self, bootstrap_servers: List[str]):
        """
        初始化Kafka工具类。

        参数：
            bootstrap_servers (List[str]): Kafka服务器地址列表。
        """
        logger.info(f"Initializing Kafka tool, Service ip: {bootstrap_servers}")
        self.bootstrap_servers = bootstrap_servers

    @retry(stop_max_attempt_number=5, wait_fixed=2000)
    def get_message_count(self, topic: str) -> int:
        """
        获取某主题下的消息总数。

        参数：
            topic (str): 主题名称。

        返回值：
            int: 主题下的消息总数。
        """
        consumer = KafkaConsumer(bootstrap_servers=self.bootstrap_servers)

        # 获取主题的所有分区
        partitions = consumer.partitions_for_topic(topic)

        if not partitions:
            logger.info(f"No partitions found for topic {topic}")
            return 0

        message_count = 0
        for partition in partitions:
            tp = TopicPartition(topic, partition)

            # 获取最早和最新的偏移量
            consumer.assign([tp])
            consumer.seek_to_beginning(tp)
            start_offset = consumer.position(tp)

            consumer.seek_to_end(tp)
            end_offset = consumer.position(tp)

            # 计算消息数目
            message_count += end_offset - start_offset

        consumer.close()
        return message_count

    @retry(stop_max_attempt_number=5, wait_fixed=2000)
    def create_producer(self) -> KafkaProducer:
        """
        创建Kafka生产者对象。

        返回值：
            KafkaProducer: Kafka生产者对象。
        """
        logger.info(f"Create Producer ing..., Service ip: {self.bootstrap_servers}")
        return KafkaProducer(
            bootstrap_servers=self.bootstrap_servers,
            value_serializer=lambda v: json.dumps(v).encode('utf-8')
        )

    @retry(stop_max_attempt_number=5, wait_fixed=2000)
    def create_consumer(self, topic: str, group_id: str, auto_offset_reset: str = 'earliest') -> KafkaConsumer:
        """
        创建Kafka消费者对象。

        参数：
            topic (str): 主题名称。
            group_id (str): 消费者组ID。
            auto_offset_reset (str): 偏移量重置策略（默认为 'earliest'）。

        返回值：
            KafkaConsumer: Kafka消费者对象。
        """
        logger.info(f"Create Consumer ing..., Service ip: {self.bootstrap_servers}")
        return KafkaConsumer(
            topic,
            bootstrap_servers=self.bootstrap_servers,
            group_id=group_id,
            auto_offset_reset=auto_offset_reset,
            value_deserializer=lambda v: json.loads(v.decode('utf-8'))
        )

    @retry(stop_max_attempt_number=5, wait_fixed=2000)
    def send_message(self, producer: KafkaProducer, topic: str, key: str, value: Dict[str, Any]) -> None:
        """
        发送消息到指定主题。

        参数：
            producer (KafkaProducer): Kafka生产者对象。
            topic (str): 主题名称。
            key (str): 消息键。
            value (Dict[str, Any]): 消息内容。
        """
        logger.info(f"Send Message ing..., Service ip: {self.bootstrap_servers}")
        future = producer.send(topic, key=key.encode('utf-8'), value=value)
        try:
            record_metadata = future.get(timeout=10)
            logger.info(
                f'[{self.bootstrap_servers}]-Message sent to {record_metadata.topic} partition {record_metadata.partition} offset {record_metadata.offset}')
        except KafkaError as e:
            logger.error(f'[{self.bootstrap_servers}]-Failed to send message: {e}')
            raise

    @retry(stop_max_attempt_number=5, wait_fixed=2000)
    def consume_messages(self, consumer: KafkaConsumer) -> Generator[Dict[str, Any], None, None]:
        """
        消费消息并返回。

        参数：
            consumer (KafkaConsumer): Kafka消费者对象。

        返回值：
            Generator[Dict[str, Any], None, None]: 生成器，返回消息字典。
        """
        logger.info(f"Consume Message ing..., Service ip: {self.bootstrap_servers}")
        try:
            for message in consumer:
                yield {
                    'key': message.key.decode('utf-8') if message.key else None,
                    'value': message.value
                }
        except KafkaError as e:
            logger.error(f'[{self.bootstrap_servers}]-Failed to consume message: {e}')
            raise


# 使用示例
if __name__ == '__main__':
    bootstrap_servers = ['192.168.1.116:9092']
    kafka_tool = KafkaTool(bootstrap_servers)

    # 创建消费者并消费消息
    consumer = kafka_tool.create_consumer('starmerx_stock_rack', 'test_group')
    try:
        for msg in kafka_tool.consume_messages(consumer):
            print(f'Received message: key={msg["key"]}, value={msg["value"]}')
    except Exception as e:
        logger.error(f"Error consuming messages: {e}")

    try:
        print(kafka_tool.get_message_count('starmerx_stock_rack'))
    except Exception as e:
        logger.error(f"Error getting message count: {e}")
